Support integer indexing of Categorical and MultivariateNormal predictions - #442

Depends on #441. This branch starts from fix/normal-fixed-params, because #441 adds the indexing tests that this PR changes. Until #441 merges, the diff here also shows its commit. Review the last commit only, or wait for #441.

dist[0] (an integer key) raised an error for Bernoulli, k_categorical(k) and MultivariateNormal(k). Distn.__getitem__ passes params[:, i] to the constructor, which is a one-dimensional array for an integer key. The constructors of these classes read params.shape[1]. Slices such as dist[1:4] worked. The other distributions in ngboost.distns accept an integer key.

Design

The fix is in the constructors. dist[i] is now the distribution of row i without the row axis, as for the other distributions and as for a numpy array:

• Categorical.__init__ builds the logits with shape (K,) + params.shape[1:], and sample1 draws probs.shape[1:] random numbers in place of (1, len(self)).
• MultivariateNormal and get_chol_factor use ... and the last axes in place of fixed axis numbers in __init__, summaries, logpdf and rv. scipy_distribution() gives one scipy object for one row and a list for a batch.
• ClassificationDistn.predict takes the argmax over the last axis. For a batch this is the same as axis 1.

I also looked at two other places for the fix. An override of __getitem__ in the three classes, or a change in Distn.__getitem__, would pass params[:, [i]] and keep the row axis. Then dist[i] would be the same as dist[i:i+1], a batch of one row, with class_probs() of shape (1, K) and a MultivariateNormal loc of shape (1, p). For Normal, dist[0].loc is a scalar, so these three classes would differ from the rest. The test from #441 compares dist[i].params with row i of dist.params, and MultivariateNormal would fail it. A change in Distn.__getitem__ would also need a check of the class, so that dist[i] stays the same for the classes that work now. The constructor fix needs no such check, and it does not touch Distn.__getitem__ or the Distn.__getattr__ that #433 adds.

For a batch, the outputs do not change. I compared the old and new code on batches of 7 rows for k_categorical(2), k_categorical(3), MultivariateNormal(2) and MultivariateNormal(3): class_probs, predict, seeded sample, loc, cov, logpdf, summaries, and the LogScore d_score and metric of the manifold are bit-identical. A batch of one row (dist[0:1]) also gives the same logpdf and seeded sample as before.

Example

d = Dist(np.random.default_rng(0).normal(scale=0.3, size=(Dist.n_params, 4))), then d[0]:

Iteration (for row in dist) now works for these classes too, because it calls __getitem__ with 0, 1, 2 and so on.

One difference stays. BetaBernoulli comes from the sympy factory, and its dist[0].class_probs() has shape (1, 2) because its class_probs calls np.column_stack. This PR does not change it.

Tests

• test_integer_index_keeps_class_and_params no longer marks Bernoulli, k_categorical and MultivariateNormal as xfail, and the unused _NO_INT_INDEX set is gone.
• test_categorical_integer_index_gives_one_row checks class_probs, predict and sample of dist[2] and dist[-1] against the batch, and iteration over the rows.
• test_multivariatenormal_integer_index_gives_one_row checks predict, cov, logpdf, scipy_distribution and sample for k = 2 and 3, and logpdf of one row at several points.
• test_integer_index_of_fitted_classifier_and_mvn checks pred_dist(X)[0] of a fitted NGBClassifier against predict_proba and predict, and of a fitted NGBRegressor with MultivariateNormal(2) against predict and logpdf.
• All 8 of these test cases fail on Store params in NormalFixedVar and NormalFixedMean so they can be sliced #441 and pass here.
• Full suite: 210 passed, 59 skipped (202 passed, 59 skipped, 3 xfailed on Store params in NormalFixedVar and NormalFixedMean so they can be sliced #441). With --slow: 244 passed, 1 skipped, 2 xfailed, 22 xpassed (236 passed, 1 skipped, 5 xfailed, 22 xpassed on Store params in NormalFixedVar and NormalFixedMean so they can be sliced #441). The other 2 xfailed and 22 xpassed cases come from the xfail mark on test_dists_runs_on_examples_logscore, which this PR does not touch (Mark only the T case as xfail in the logscore distribution test #437 changes it). black, isort, flake8 and pylint are clean on the changed files.

🤖 Generated with Claude Code